import os
import asyncio
from typing import Any
from loguru import logger

from ..models.protobuf.douyin import *
from ...utils.db import MongoDBManager
from .room_fetcher import RoomInfoFetcher
from ...utils.proxy_manager import GlobalProxyManager
from ..handlers.message_handler import DouyinMessageHandler
from ..handlers.connection_manager import WebSocketConnectionManager
from .task_manager import TaskManager
from ..config.live_config import LiveConfig


class DouyinLiveStream:
    def __init__(self, shared_proxy: Any) -> None:
        """
        初始化抖音直播流处理器

        Args:
            shared_proxy: 共享的代理实例
        """
        self.collection_name: str = 'tk_forerver_urls'
        self.black_list: list = []  # 黑名单链接列表
        self.task_manager = TaskManager()
        self.config = LiveConfig()
        self.db = MongoDBManager()
        self.proxy_manager = GlobalProxyManager(shared_proxy)
        self.message_handler = DouyinMessageHandler()
        self.connection_manager = WebSocketConnectionManager(
            self.message_handler,
            self.proxy_manager
        )
        self.init_black_list()

    def init_black_list(self) -> None:
        """初始化黑名单列表"""
        urls = self.db.get_db_black_urls()
        if urls:
            self.black_list = urls

    async def close_websocket(self,
                              ws: Any,
                              url: str,
                              shared_list: list,
                              message: str) -> None:
        """
        关闭WebSocket连接

        Args:
            ws: WebSocket连接实例
            url: 直播间URL
            shared_list: 共享列表
            message: 关闭原因
        """
        pid = os.getpid()
        # logger.info(f"进程 {pid} {url} 已关闭链接 message:{message}")
        if url in shared_list:
            shared_list.remove(url)
        self.task_manager.remove_task(url)
        if ws:
            await ws.close()

    async def start_websocket(self,
                              url: str,
                              shared_queue: Any,
                              shared_list: list) -> None:
        """启动WebSocket连接"""
        try:
            fetcher = RoomInfoFetcher()
            response = await asyncio.to_thread(fetcher.fetch_live_room_info, url)

            if response:
                wss_url, ttwid, web_url = response
                await self.connection_manager.create_connection(
                    wss_url, ttwid, web_url,
                    shared_queue, shared_list,
                    self.close_websocket
                )
            else:
                logger.warning(f'{url} 获取房间信息失败')
                if url in shared_list:
                    shared_list.remove(url)
                if url in self.task_manager.task_urls:
                    self.task_manager.remove_task(url)
        except Exception as e:
            logger.error(f"启动WebSocket失败: {url} - {str(e)}")
            if url in shared_list:
                shared_list.remove(url)
            if url in self.task_manager.task_urls:
                self.task_manager.remove_task(url)

    async def run(self,
                  q: Any,
                  shared_queue: Any,
                  shared_list: list) -> None:
        """运行直播流处理"""
        async def handle_url(url: str) -> None:
            try:
                # logger.info(f'进程 {os.getpid()} 开始处理任务: {url}')
                await self.start_websocket(url, shared_queue, shared_list)
            except Exception as e:
                logger.error(f"进程 {os.getpid()} 处理任务失败 {url}: {str(e)}")
                if url in shared_list:
                    shared_list.remove(url)
                if url in self.task_manager.task_urls:
                    self.task_manager.remove_task(url)

        while True:  # 添加循环来持续检查队列
            try:
                # 检查队列状态
                if not q.empty():
                    queue_size = q.qsize()
                    # logger.info(f"进程 {os.getpid()} 队列中有任务,当前任务数: {queue_size}")
                    
                    # 处理队列任务
                    await self.task_manager.process_queue(
                        q, shared_list, handle_url
                    )
                
                # # 定期输出状态信息
                # logger.info(
                #     f"进程 {os.getpid()} 当前监控任务数: {len(self.task_manager.task_urls)}/"
                #     f"{self.config.MAX_TASK}, 队列剩余: {q.qsize()}"
                # )

                await asyncio.sleep(self.config.SLEEP_INTERVAL)

            except Exception as e:
                logger.error(f"处理任务出错: {str(e)}")
                await asyncio.sleep(self.config.SLEEP_INTERVAL)

    def get_black_urls(self) -> list:
        """获取黑名单列表"""
        return self.black_list

    def get_db_urls(self) -> list:
        """获取数据库中存储的url列表"""
        return self.db.find_all(
            collection_name=self.collection_name,
            attribute='url'
        )
